from llvmlite import binding as llvm
from ctypes import Structure, c_int32, c_double, POINTER, CFUNCTYPE

# 初始化 LLVM
llvm.initialize()
llvm.initialize_native_target()
llvm.initialize_native_asmprinter()

# 定义 LLVM 汇编代码（包含自定义结构体操作）
ir_text = """
%MyStruct = type { i32, double }

; 返回结构体的函数
define %MyStruct @create_struct() {
  ret %MyStruct { i32 42, double 3.140000e+00 }
}

; 接受结构体作为参数的函数
define double @get_double_field(%MyStruct %s) {
  %field = extractvalue %MyStruct %s, 1
  ret double %field
}

; 通过指针修改结构体的函数（推荐方式）
define void @modify_struct(%MyStruct* %ptr) {
  %field_ptr = getelementptr inbounds %MyStruct, %MyStruct* %ptr, i32 0, i32 0
  store i32 100, i32* %field_ptr
  ret void
}

; 创建结构体并返回指针的函数
define %MyStruct* @create_struct_ptr() {
  %ptr = alloca %MyStruct
  store %MyStruct { i32 200, double 6.280000e+00 }, %MyStruct* %ptr
  ret %MyStruct* %ptr
}
"""

# 解析汇编代码
module = llvm.parse_assembly(ir_text)
module.verify()  # 验证模块的有效性

# 创建执行引擎
target = llvm.Target.from_default_triple()
target_machine = target.create_target_machine()
engine = llvm.create_mcjit_compiler(module, target_machine)

# 定义匹配的 Python 结构体（必须与 LLVM 内存布局一致）
class MyStruct(Structure):
    _fields_ = [("a", c_int32), ("b", c_double)]

# 获取函数地址并设置调用签名
# 1. 返回结构体的函数
create_struct_addr = engine.get_function_address("create_struct")
create_struct = CFUNCTYPE(MyStruct)(create_struct_addr)

# 2. 接受结构体参数的函数
get_double_field_addr = engine.get_function_address("get_double_field")
get_double_field = CFUNCTYPE(c_double, MyStruct)(get_double_field_addr)

# 3. 通过指针修改结构体的函数
modify_struct_addr = engine.get_function_address("modify_struct")
modify_struct = CFUNCTYPE(None, POINTER(MyStruct))(modify_struct_addr)

# 4. 返回结构体指针的函数
create_struct_ptr_addr = engine.get_function_address("create_struct_ptr")
create_struct_ptr = CFUNCTYPE(POINTER(MyStruct))(create_struct_ptr_addr)

# 测试函数调用
if __name__ == "__main__":
    print("=== 直接返回结构体 ===")
    result = create_struct()
    print(f"create_struct: a={result.a}, b={result.b}")  # 输出: a=42, b=3.14
    
    print("\n=== 结构体作为参数 ===")
    s = MyStruct(10, 2.718)
    d = get_double_field(s)
    print(f"get_double_field: {d}")  # 输出: 2.718
    
    print("\n=== 通过指针修改结构体 ===")
    s_ptr = MyStruct(1, 1.0)
    print(f"修改前: a={s_ptr.a}")
    modify_struct(s_ptr)
    print(f"修改后: a={s_ptr.a}")  # 输出: 100
    
    print("\n=== 返回结构体指针 ===")
    ptr = create_struct_ptr()
    print(f"指针内容: a={ptr.contents.a}, b={ptr.contents.b}")  # 输出: a=200, b=6.28
